import pandas as pd
import numpy as np
from datetime import datetime

class Strategy:
    def judge2(self,index_data_df,dt,delay=20):
        """
        传入指数数据和日期，返回当天的目标持仓
        :param index_data_df: 指数的数据
        :param dt: 当天日期
        :param delay: 与多少天前的收盘价比较
        :return: 当日的目标持仓
        """
        if type(dt) != type('123'):
            dt = datetime.strftime(dt,'%Y-%m-%d')
        dt_loc = list(index_data_df.index).index(dt)
        if dt_loc<delay:
            # 直接跳过
            print(dt + ' has no pre_dt')
            return None
        else:
            # 找出N日涨幅最大的指数，第二天持有
            pre_dt = index_data_df.index[dt_loc-delay]
            now_df = index_data_df.loc[dt,:]
            pre_df = index_data_df.loc[pre_dt,:]
            max_increase = np.max(now_df / pre_df)
            if (max_increase<=1):
                # 说明各个指数的N天收益率均为负，此时选择空仓
                return None
            else:
                buy = index_data_df.columns[np.argmax(now_df/pre_df)]
                return buy
            # # 找出大于20日均线最大的那个指数，买入
            # mean_df = index_data_df.iloc[dt_loc-delay:dt_loc,:].mean()
            # now_df = index_data_df.iloc[dt_loc, :]
            # max_increase = np.max(now_df / mean_df)
            # if max_increase<=1:
            #     # 说明各个指数的N天收益率均为负，此时选择空仓
            #     return None
            # else:
            #     buy = index_data_df.columns[np.argmax(now_df/mean_df)]
            #     return buy

    def is_stop_loss(self,context,invest_code,index_etf_data_df,dt,stop_loss_rate=0.1):
        # 限价止损函数，判断当前持仓是否触发止损
        buy_price = context.now_have[invest_code]['price']
        now_price = index_etf_data_df.loc[dt,invest_code]
        if now_price <= buy_price * (1-stop_loss_rate):
            return True
        else:
            return False

    def is_stop_win(self,context,invest_code,index_etf_data_df,dt,stop_win_rate=0.05):
        # 跟踪止盈函数，判断当前持仓是否触发止盈
        buy_dt = context.now_have[invest_code]['dt']
        highest_price = np.max(index_etf_data_df.loc[buy_dt:dt,invest_code])
        now_price = index_etf_data_df.loc[dt, invest_code]
        if now_price <= highest_price * (1-stop_win_rate):
            return True
        else:
            return False